import numpy as np
from rouge import Rouge


class RougeScore:
    def __init__(self, tokenizer):
        self.rouge = Rouge()
        self.tokenizer = tokenizer
        self.metrics = np.zeros(3)
        self.n = 0

    def decode(self, ids):
        ids = ids[:, 1:].cpu().numpy()
        texts = self.tokenizer.batch_decode(ids, skip_special_tokens=True)
        texts = [s.replace(' ', '') for s in texts]
        return texts

    def update(self, yt, yp):
        yt = self.decode(yt)
        yp = self.decode(yp)
        for t, p, in zip(yt, yp):
            try:
                scores = self.rouge.get_scores(hyps=t, refs=p)[0]
                scores = [scores['rouge-1']['f'], scores['rouge-2']['f'], scores['rouge-l']['f']]
            except ValueError:
                scores = [0, 0, 0]
            self.metrics += np.array(scores)
            self.n += 1

    def accumulate(self, weight=None):
        if weight is None:
            ret = self.metrics / self.n
        else:
            ret = self.metrics.dot(weight) / self.n
        self.reset()
        return ret

    def reset(self):
        self.metrics = np.zeros(3)
        self.n = 0
